import tensorflow as tf
import numpy as np
from    tensorflow.keras import layers, Sequential

class DispNet(tf.keras.Model):
    def __init__(self):
        super(DispNet, self).__init__()
        self.conv1 = layers.Conv2D(64,(7,7),padding = 'same',activation = 'relu')
        self.pool1 = layers.MaxPool2D((2,2))
        self.conv3 = layers.Conv2D(128,(5,5),padding = 'same',activation = 'relu')
        self.pool3 = layers.MaxPool2D((2,2))
        self.conv17 = layers.Conv2D(256,(5,5),padding = 'same',activation = 'relu')
        self.pool8 = layers.MaxPool2D((2,2))
        self.conv4 = layers.Conv2D(256,(3,3),padding = 'same',activation = 'relu')
        self.conv9 = layers.Conv2D(512,(3,3),padding = 'same',activation = 'relu')
        self.pool5 = layers.MaxPool2D((2,2))
        self.conv10 = layers.Conv2D(512,(3,3),padding = 'same',activation = 'relu')
        self.conv11 = layers.Conv2D(512,(3,3),padding = 'same',activation = 'relu')
        self.pool6 = layers.MaxPool2D((2,2))
        self.conv12 = layers.Conv2D(512,(3,3),padding = 'same',activation = 'relu')
        self.conv13 = layers.Conv2D(1024,(3,3),padding = 'same',activation = 'relu')
        self.pool7 = layers.MaxPool2D((2,2))
        self.conv14 = layers.Conv2D(1024,(3,3),padding = 'same',activation = 'relu')
        self.conv18 = layers.Conv2D(1,(3,3),padding = 'same',activation = 'relu')
        self.up1 = layers.UpSampling2D((2,2))
        self.deconv4 = layers.Conv2DTranspose(512,(4,4),strides=(2, 2),padding = 'same',activation = 'relu')
        self.bn1 = layers.BatchNormalization()
        self.conv19 = layers.Conv2D(512,(3,3),padding = 'same',activation = 'relu')
        self.conv20 = layers.Conv2D(1,(3,3),padding = 'same',activation = 'relu')
        self.up2 = layers.UpSampling2D((2,2))
        self.deconv5 = layers.Conv2DTranspose(256,(4,4),strides=(2, 2),padding = 'same',activation = 'relu')
        self.bn2 = layers.BatchNormalization()
        self.conv21 = layers.Conv2D(256,(3,3),padding = 'same',activation = 'relu')
        self.deconv24 = layers.Conv2DTranspose(128,(4,4),strides=(2, 2),padding = 'same',activation = 'relu')
        self.bn3 = layers.BatchNormalization()
        self.conv22 = layers.Conv2D(1,(3,3),padding = 'same',activation = 'relu')
        self.up3 = layers.UpSampling2D((2,2))
        self.conv23 = layers.Conv2D(128,(3,3),padding = 'same',activation = 'relu')
        self.conv24 = layers.Conv2D(1,(3,3),padding = 'same',activation = 'relu')
        self.up4 = layers.UpSampling2D((2,2))
        self.deconv7 = layers.Conv2DTranspose(64,(4,4),strides=(2, 2),padding = 'same',activation = 'relu')
        self.bn4 = layers.BatchNormalization()
        self.conv25 = layers.Conv2D(64,(3,3),padding = 'same',activation = 'relu')
        self.conv26 = layers.Conv2D(1,(3,3),padding = 'same',activation = 'relu')
        self.up5 = layers.UpSampling2D((2,2))
        self.deconv8 = layers.Conv2DTranspose(32,(4,4),strides=(2, 2),padding = 'same',activation = 'relu')
        self.bn5 = layers.BatchNormalization()
        self.conv27 = layers.Conv2D(32,(3,3),padding = 'same',activation = 'relu')
        self.conv28 = layers.Conv2D(1,(3,3),padding = 'same',activation = 'relu')
    def call(self, inputs, training=None):
        c1 = self.conv1(inputs)
        p1 = self.pool1(c1)
        c3 = self.conv3(p1)
        p3 = self.pool3(c3)
        c17 = self.conv17(p3)
        p8 = self.pool8(c17)
        c4 = self.conv4(p8)
        c9 = self.conv9(c4)
        p5 = self.pool5(c9)
        c10 = self.conv10(p5)
        c11 = self.conv11(c10)
        p6 = self.pool6(c11)
        c12 = self.conv12(p6)
        c13 = self.conv13(c12)
        p7 = self.pool7(c13)
        c14 = self.conv14(p7)
        c18 = self.conv18(c14)
        u1 = self.up1(c18)
        d4 = self.deconv4(c14)
        b1 = self.bn1(d4)
        merge_2 = tf.concat([c12,b1,u1],axis = 3)
        c19 = self.conv19(merge_2)
        c20 = self.conv20(c19)
        u2 = self.up2(c20)
        d5 = self.deconv5(c19)
        b2 = self.bn2(d5)
        merge_3 = tf.concat([c10,b2,u2],axis = 3)
        c21 = self.conv21(merge_3)
        d24 = self.deconv24(c21)
        b3 = self.bn3(d24)
        c22 = self.conv22(c21)
        u3 = self.up3(c22)
        merge_4 = tf.concat([c4,b3,u3],axis = 3)
        c23 = self.conv23(merge_4)
        c24 = self.conv24(c23)
        u4 = self.up4(c24)
        d7 = self.deconv7(c23)
        b4 = self.bn4(d7)
        merge_5 = tf.concat([p3,b4,u4],axis = 3)
        c25 = self.conv25(merge_5)
        c26 = self.conv26(c25)
        u5 = self.up5(c26)
        d8 = self.deconv8(c25)
        b5 = self.bn5(d8)
        merge_6 = tf.concat([p1,b5,u5],axis = 3)
        c27 = self.conv27(merge_6)
        out = self.conv28(c27)
        return out,c26,c24,c22,c20
